package com.xjai.controller;


import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.stream.Collectors;

@Slf4j
@RestController
@RequestMapping("ollama")
public class OllamaController {

    @Resource
    private OllamaChatClient ollamaChatClient;

    @GetMapping("/ai/chat")
    public Object aiOllamaCall(@RequestParam String msg){
       log.info("prompt ------->> 111 :" + msg);
       return ollamaChatClient.call(msg);
    }

    @GetMapping("/ai/stream1")
    public Flux<ChatResponse> aiOllamaStream1(@RequestParam String msg){
        Prompt prompt = new Prompt(new UserMessage(msg));
        Flux<ChatResponse> stream = ollamaChatClient.stream(prompt);
        return stream;
    }


    @GetMapping("/ai/stream2")
    public List<String> aiOllamaStream2(@RequestParam String msg){
        Prompt prompt = new Prompt(new UserMessage(msg));
        Flux<ChatResponse> stream = ollamaChatClient.stream(prompt);

        List<String> collect = stream.toStream().map(chatResponse -> {
            String content = chatResponse.getResult().getOutput().getContent();
            log.info("content" + content);
            return content;
        }).collect(Collectors.toList());

        return collect;
    }





}
